package Share;

import java.math.BigInteger;
import java.security.AlgorithmParameterGenerator;
import java.security.AlgorithmParameters;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.X509EncodedKeySpec;
import java.util.Formatter;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.spec.DHParameterSpec;



/**
 * Main class used for symetrical cryptographic operations (DH Key agreement and 3DES ciphering).
 * 3DES is used to generate a keystream that will be xored with the bytes to cipher. 
 * To do that, an IV is passed to the crypto function, then the IV get ciphered, and used as the xor-key.
 * When all the bytes are used, the IV is ciphered again.
 * IV should never be reused, that's why we are using a SecureRadom generator. 
 *
 */
public class SymCrypto {

	//Code inspired by : http://www.exampledepot.com/egs/javax.crypto/GenDhParams.html
	// 					 http://www.exampledepot.com/egs/javax.crypto/DesString.html
	//
	private BigInteger _P;
	private BigInteger _G;
	private int _L;
	private   DHParameterSpec _dhSpec;
	private  KeyPair _keypair;
	private SecretKey _sharedKey;
	private Cipher _eCipher;
	private Cipher _dCipher;
	private boolean _isReady;
	private byte[]  _IV;  //Its size should be a multiple of 8 bytes
	private int    _remainingBytes=-1;
	private SecureRandom _random;
	
	public static int IVLEN=16;

	//----------------------------------------
	// Constructors
	//----------------------------------------
	
	/**
	 * Class constructor for the client (who generates the DH params)
	 */
	public SymCrypto() {
		genDhParams();
		_random = new SecureRandom();
		_IV= new byte[IVLEN];
		 

	}
	/**
	 * Class constructor for the server (who receives the DH parameters from the client
	 * @param P String representation of the BigInteger P
	 * @param G String representation of the BigInteger G
	 * @param L String representation of the int l
	 * @param B64PublicKey The g^a [p] from the client, in a base 64 representation
	 */
	public SymCrypto(String P, String G, String L, byte[] publicKeyBytes) {
		_P = new BigInteger(P);
		_G = new BigInteger(G);
		_L=  Integer.parseInt(L);
		_dhSpec = new DHParameterSpec(_P, _G, _L);
		_random = new SecureRandom();
		_IV= new byte[IVLEN];
		
		genKeyPair();
		genSharedKeyAndCiphers(publicKeyBytes);
		_IV = new byte[16]; 
	}
	
	//---------------------------------------------------------------------------------------
	// Generation of the parameters used by the crypto system.
	//---------------------------------------------------------------------------------------
	
	private void genDhParams() {
		try {
			System.out.println("Generating the numbers P and G");
			// Create the parameter generator for a 1024-bit DH key pair
			AlgorithmParameterGenerator paramGen = AlgorithmParameterGenerator.getInstance("DH");
			paramGen.init(1024);

			// Generate the parameters
			AlgorithmParameters params = paramGen.generateParameters();
			_dhSpec
			= (DHParameterSpec)params.getParameterSpec(DHParameterSpec.class);

			// Return the three values in a string
			_P =  _dhSpec.getP();
			_G =_dhSpec.getG();
			_L = _dhSpec.getL();
			genKeyPair();
		}
		catch (Exception e) {
			System.out.println("Error while generating DH values");
			System.exit(-1);
		}


	}
	private void genKeyPair() {
		try {
			KeyPairGenerator keyGen;
			System.out.println("Generating the KeyPair");
			keyGen = KeyPairGenerator.getInstance("DH");
			keyGen.initialize(_dhSpec);
			_keypair = keyGen.generateKeyPair();
		}
		catch (Exception e) {
			System.out.println("Error while generating DH values");
			System.exit(-1);
		}
	}
	public void genSharedKeyAndCiphers(byte[] publicKeyBytes ) {
		try {

		
	    // Convert the public key bytes into a PublicKey object
	    X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(publicKeyBytes);
	    KeyFactory keyFact = KeyFactory.getInstance("DH");
	    PublicKey publicKey = keyFact.generatePublic(x509KeySpec);

	    // Prepare to generate the secret key with the private key and public key of the other party
	    KeyAgreement ka = KeyAgreement.getInstance("DH");
	    ka.init(_keypair.getPrivate());
	    ka.doPhase(publicKey, true);

	    // Specify the type of key to generate;
	    // see Listing All Available Symmetric Key Generators
	    String algorithm = "DESede";

	    // Generate the secret key
	    System.out.println("Calculating the shared key");
	    _sharedKey = ka.generateSecret(algorithm);
	    _eCipher = Cipher.getInstance(algorithm);
	    _dCipher = Cipher.getInstance(algorithm);
        _eCipher.init(Cipher.ENCRYPT_MODE, _sharedKey);
        _dCipher.init(Cipher.DECRYPT_MODE, _sharedKey); //Useless here
        
        System.out.println( bytesToHexString(_sharedKey.getEncoded()));
        
		}
		catch (Exception e) {
			System.out.println("Error while generating the private key");
			System.exit(-1);
		}
		
	}
	
	//----------------------------------------
	// 3DES Methods
	//----------------------------------------
		
		public void setNewIV(byte[] IV) {
			
			this._IV = IV;
			_remainingBytes = -1;
			_isReady =true;
			System.out.println(bytesToHexString(_IV)) ;
		}
		
		public byte[] genNewIV() {
			
		 _random.nextBytes(_IV);
		 _remainingBytes=-1;
		_isReady =true;
		System.out.println(bytesToHexString(_IV)) ;
		 return _IV;
		 
		}
		

		
		
		public int encrypt(byte b) {
	        
			if (_remainingBytes < 1) {
				try {
	
		           _IV = _eCipher.doFinal(_IV);
		            _remainingBytes = _IV.length;
		           
		           
	
		            
		            
		        } catch (Exception e) {
		        	System.out.println("Error while encrypting");
		        	return -1;
		        }
			}
			
			return b ^  _IV[--_remainingBytes];
	        
	    }
		public int decrypt(byte b) {
			if (_remainingBytes < 1) {
				try {
	
		           _IV = _eCipher.doFinal(_IV);
		            _remainingBytes = _IV.length;
		            
	
		            
		            
		        } catch (Exception e) {
		        	System.out.println("Error while encrypting");
		        	return -1;
		        }
			}
			 
			return b  ^ _IV[--_remainingBytes];
	        
	    }


	//---------------------------------------
	//  Getters
	//---------------------------------------
	public BigInteger getP() {
		return _P;
	}
	public BigInteger getG() {
		return _G;
	}
	public int getL() {
		return _L;
	}

	public byte[] getPublicKeyBytes() {

		return _keypair.getPublic().getEncoded();

	}
	/**
	 * Return true when the Crypto system is ready (ie the DH shared key is computed, and 3DES ciphers are ready, and IV is set)
	 * @return
	 */
	public boolean isReady() {
		return _isReady;
	}
	
	//Used for debug purpose
	public static String bytesToHexString(byte[] bytes) {
	    StringBuilder sb = new StringBuilder(bytes.length * 2);

	    Formatter formatter = new Formatter(sb);
	    for (byte b : bytes) {
	        formatter.format("%02x", b);
	    }

	    return sb.toString();
	}
	
}
